1use crate::ext::io::*;
3use crate::ext::psb::*;
4use crate::scripts::base::*;
5use crate::types::*;
6use crate::utils::encoding::*;
7use crate::utils::img::*;
8use anyhow::Result;
9use emote_psb::PsbReader;
10use libtlg_rs::TlgColorType;
11use std::collections::HashMap;
12use std::io::Read;
13use std::path::{Path, PathBuf};
14use url::Url;
15
16#[derive(Debug)]
17pub struct DrefBuilder {}
19
20impl DrefBuilder {
21 pub fn new() -> Self {
23 Self {}
24 }
25}
26
27impl ScriptBuilder for DrefBuilder {
28 fn default_encoding(&self) -> Encoding {
29 Encoding::Cp932
30 }
31
32 fn build_script(
33 &self,
34 buf: Vec<u8>,
35 filename: &str,
36 encoding: Encoding,
37 _archive_encoding: Encoding,
38 config: &ExtraConfig,
39 archive: Option<&Box<dyn Script>>,
40 ) -> Result<Box<dyn Script>> {
41 Ok(Box::new(Dref::new(
42 buf, encoding, filename, config, archive,
43 )?))
44 }
45
46 fn extensions(&self) -> &'static [&'static str] {
47 &["dref"]
48 }
49
50 fn script_type(&self) -> &'static ScriptType {
51 &ScriptType::EmoteDref
52 }
53
54 fn is_image(&self) -> bool {
55 true
56 }
57}
58
59struct Dpak {
60 psb: VirtualPsbFixed,
61}
62
63struct OffsetData {
64 left: u32,
65 top: u32,
66}
67
68impl Dpak {
69 pub fn new<P: AsRef<Path>>(path: P) -> Result<Self> {
70 let f = std::fs::File::open(path)?;
71 let mut f = std::io::BufReader::new(f);
72 let mut psb = PsbReader::open_psb(&mut f)
73 .map_err(|e| anyhow::anyhow!("Failed to read PSB from DPAK: {:?}", e))?;
74 let psb = psb
75 .load()
76 .map_err(|e| anyhow::anyhow!("Failed to load PSB from DPAK: {:?}", e))?;
77 let psb = psb.to_psb_fixed();
78 Ok(Self { psb })
79 }
80
81 pub fn load_from_data(data: &[u8]) -> Result<Self> {
82 let mut psb = PsbReader::open_psb(MemReaderRef::new(data))
83 .map_err(|e| anyhow::anyhow!("Failed to read PSB from DPAK data: {:?}", e))?;
84 let psb = psb
85 .load()
86 .map_err(|e| anyhow::anyhow!("Failed to load PSB from DPAK data: {:?}", e))?;
87 let psb = psb.to_psb_fixed();
88 Ok(Self { psb })
89 }
90
91 pub fn load_image(&self, name: &str) -> Result<(ImageData, Option<OffsetData>)> {
92 let root = self.psb.root();
93 let rid = root[name]
94 .resource_id()
95 .ok_or_else(|| anyhow::anyhow!("Resource ID for image '{}' not found in DPAK", name))?
96 as usize;
97 if rid >= self.psb.resources().len() {
98 return Err(anyhow::anyhow!(
99 "Resource ID {} out of bounds for DPAK with {} resources",
100 rid,
101 self.psb.resources().len()
102 ));
103 }
104 let resource = &self.psb.resources()[rid];
105 Self::load_img(&resource)
106 }
107
108 fn load_img(data: &[u8]) -> Result<(ImageData, Option<OffsetData>)> {
109 if libtlg_rs::is_valid_tlg(data) {
110 Ok((Self::load_tlg(data)?, None))
111 } else {
112 Self::load_png(data)
113 }
114 }
115
116 fn load_tlg(data: &[u8]) -> Result<ImageData> {
117 let img = libtlg_rs::load_tlg(MemReaderRef::new(data))
118 .map_err(|e| anyhow::anyhow!("Failed to decode TLG image: {:?}", e))?;
119 let color = img.color;
120 let mut re = ImageData {
121 width: img.width as u32,
122 height: img.height as u32,
123 color_type: match img.color {
124 TlgColorType::Grayscale8 => ImageColorType::Grayscale,
125 TlgColorType::Bgr24 => ImageColorType::Bgr,
126 TlgColorType::Bgra32 => ImageColorType::Bgra,
127 },
128 data: img.data,
129 depth: 8,
130 };
131 if let Some(v) = img.tags.get(&Vec::from(b"mode")) {
132 if v == b"alpha" && color == TlgColorType::Bgr24 {
133 convert_bgr_to_bgra(&mut re)?;
134 }
135 }
136 Ok(re)
137 }
138
139 fn load_png(data: &[u8]) -> Result<(ImageData, Option<OffsetData>)> {
140 let mut img = load_png(MemReaderRef::new(&data))?;
141 match img.color_type {
142 ImageColorType::Rgb => {
143 convert_rgb_to_rgba(&mut img)?;
144 }
145 _ => {}
146 }
147 Ok((
148 img,
149 Self::try_read_offset_from_png(MemReaderRef::new(&data))?,
150 ))
151 }
152
153 fn try_read_offset_from_png(mut data: MemReaderRef) -> Result<Option<OffsetData>> {
154 data.pos = 8; data.pos += 8; data.pos += 17; loop {
158 let chunk_size = data.read_u32_be()?;
159 let mut chunk_type = [0u8; 4];
160 data.read_exact(&mut chunk_type)?;
161 if &chunk_type == b"IDAT" || &chunk_type == b"IEND" {
162 break;
163 }
164 if &chunk_type == b"oFFs" {
165 let x = data.read_u32_be()?;
166 let y = data.read_u32_be()?;
167 if data.read_u8()? == 0 {
168 return Ok(Some(OffsetData { left: x, top: y }));
169 }
170 }
171 data.pos += chunk_size as usize + 4; }
173 Ok(None)
174 }
175}
176
177#[derive(Default)]
178struct DpakLoader {
179 map: HashMap<String, Dpak>,
180}
181
182impl DpakLoader {
183 pub fn load_image(
184 &mut self,
185 dir: &Path,
186 dpak: &str,
187 filename: &str,
188 ) -> Result<(ImageData, Option<OffsetData>)> {
189 let dpak = match self.map.get(dpak) {
190 Some(d) => d,
191 None => {
192 let path = dir.join(dpak);
193 let ndpak = Dpak::new(&path)?;
194 self.map.insert(dpak.to_string(), ndpak);
195 self.map.get(dpak).unwrap()
196 }
197 };
198 dpak.load_image(filename)
199 }
200
201 pub fn load_archives(&mut self, in_archives: &HashMap<String, Vec<u8>>) -> Result<()> {
202 for (name, data) in in_archives.iter() {
203 if !self.map.contains_key(name) {
204 let dpak = Dpak::load_from_data(data)?;
205 self.map.insert(name.clone(), dpak);
206 }
207 }
208 Ok(())
209 }
210}
211
212pub struct Dref {
214 urls: Vec<Url>,
215 dir: PathBuf,
216 in_archives: HashMap<String, Vec<u8>>,
217}
218
219impl std::fmt::Debug for Dref {
220 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
221 f.debug_struct("Dref")
222 .field("urls", &self.urls)
223 .field("dir", &self.dir)
224 .finish()
225 }
226}
227
228impl Dref {
229 pub fn new(
237 buf: Vec<u8>,
238 encoding: Encoding,
239 filename: &str,
240 _config: &ExtraConfig,
241 archive: Option<&Box<dyn Script>>,
242 ) -> Result<Self> {
243 let text = decode_with_bom_detect(encoding, &buf, true)?.0;
244 let mut urls = Vec::new();
245 for text in text.lines() {
246 let text = text.trim();
247 if text.is_empty() {
248 continue;
249 }
250 urls.push(Url::parse(text)?);
251 }
252 let path = Path::new(filename);
253 let dir = if let Some(parent) = path.parent() {
254 parent.to_path_buf()
255 } else {
256 PathBuf::from(".")
257 };
258 if urls.is_empty() {
259 return Err(anyhow::anyhow!("No URLs found in DREF file: {}", filename));
260 }
261 for u in urls.iter() {
262 if u.scheme() != "psb" {
263 return Err(anyhow::anyhow!(
264 "Invalid URL scheme in DREF file: {} (expected 'psb')",
265 u
266 ));
267 }
268 }
269 let mut in_archives = HashMap::new();
270 if let Some(archive) = archive {
271 if archive.is_archive() {
272 for url in urls.iter() {
273 let filename = url.domain().ok_or(anyhow::anyhow!(
274 "Invalid URL in DREF file: {} (missing domain)",
275 url
276 ))?;
277 if let Ok(mut content) = archive.open_file_by_name(filename, true) {
278 in_archives.insert(filename.to_string(), content.data()?);
279 }
280 }
281 }
282 }
283 Ok(Self {
284 urls,
285 dir,
286 in_archives,
287 })
288 }
289}
290
291impl Script for Dref {
292 fn default_output_script_type(&self) -> OutputScriptType {
293 OutputScriptType::Json
294 }
295
296 fn default_format_type(&self) -> FormatOptions {
297 FormatOptions::None
298 }
299
300 fn is_image(&self) -> bool {
301 true
302 }
303
304 fn export_image(&self) -> Result<ImageData> {
305 let mut loader = DpakLoader::default();
306 loader.load_archives(&self.in_archives)?;
307 let base_url = &self.urls[0];
308 let dpak = base_url.domain().ok_or(anyhow::anyhow!(
309 "Invalid URL in DREF file: {} (missing domain)",
310 base_url
311 ))?;
312 let (mut base_img, base_offset) =
313 loader.load_image(&self.dir, dpak, base_url.path().trim_start_matches("/"))?;
314 if let Some(o) = base_offset {
315 eprintln!("WARN: Base image offset: left={}, top={}", o.left, o.top);
316 crate::COUNTER.inc_warning();
317 }
318 for url in &self.urls[1..] {
319 let dpak = url.domain().ok_or(anyhow::anyhow!(
320 "Invalid URL in DREF file: {} (missing domain)",
321 url
322 ))?;
323 let (mut img, img_offset) =
324 loader.load_image(&self.dir, dpak, url.path().trim_start_matches("/"))?;
325 let (top, left) = match img_offset {
326 Some(o) => (o.top, o.left),
327 None => (0, 0),
328 };
329 if base_img.color_type != img.color_type {
330 if base_img.color_type == ImageColorType::Rgba
331 && img.color_type == ImageColorType::Rgb
332 {
333 convert_rgb_to_rgba(&mut img)?;
334 } else if base_img.color_type == ImageColorType::Bgra
335 && img.color_type == ImageColorType::Bgr
336 {
337 convert_bgr_to_bgra(&mut img)?;
338 } else if base_img.color_type == ImageColorType::Rgba
339 && img.color_type == ImageColorType::Bgra
340 {
341 convert_bgra_to_rgba(&mut img)?;
342 } else if base_img.color_type == ImageColorType::Bgra
343 && img.color_type == ImageColorType::Rgba
344 {
345 convert_rgba_to_bgra(&mut img)?;
346 }
347 }
348 draw_on_img_with_opacity(&mut base_img, &img, left, top, 0xff)?;
349 }
350 Ok(base_img)
351 }
352}